import imp
from matplotlib import pyplot as plt
import numpy as np
import math
from mpl_toolkits.mplot3d import Axes3D
import os


def plot2d(dir):
    out_dir  = './imgs/' + dir[dir.find('out/')+4:];
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    # fig = plt.figure(figsize=(8, 8), facecolor='w')
    fig = plt.figure()

    s128 = np.loadtxt(dir+'128.txt')
    norm = np.loadtxt(dir+'norm.txt')
    
    ax128 = Axes3D(fig)
    ax128.plot_surface(s128[0].reshape((129,129)),s128[1].reshape((129,129)),s128[2].reshape((129,129)),cmap=plt.get_cmap('rainbow'))
    # plt.title('n = 128')
    # plt.show()
    fig.savefig(out_dir+'solution.pdf')
    
    fig_norm = plt.figure()
    xtics = [4,5,6,7]
    plt.plot(xtics,norm[0],label = '1-norm')
    plt.plot(xtics, norm[1],label='2-norm')
    plt.plot(xtics ,norm[2],label = '$\infinity$-norm')
    plt.legend(loc = 'best')
    plt.xticks(xtics,['2^4','2^5','2^6','2^7'])
    plt.title('Norm')
    # plt.show()
    fig_norm.savefig(out_dir+'norm.pdf')
    
def plot1d(dir,type):
    out_dir  = './imgs/' + dir[dir.find('out/')+4:];
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    s256 = np.loadtxt(dir+'256.txt')
    norm = np.loadtxt(dir+'norm.txt')
    
    fig = plt.figure()
    plt.plot(s256[0],s256[1])
    
    if type == 1:
        plt.plot(s256[0],np.exp(s256[0]))
        function = 'exp(x)'
    elif type == 2:
        plt.plot(s256[0],s256[0]**2)
        function = 'x^2'
    elif type == 3:
        plt.plot(s256[0],np.sin(math.pi*s256[0]))
        function = 'sin(\pi x)'
    plt.legend(['my solution',function])
    plt.show()
    fig.savefig(out_dir+'solution.pdf')
    
    fig_norm = plt.figure()
    xtics = [4,5,6,7]
    plt.plot(xtics,norm[0],label = '1-norm')
    plt.plot(xtics, norm[1],label='2-norm')
    plt.plot(xtics ,norm[2],label = '$\infinity$-norm')
    plt.legend(loc = 'best')
    plt.xticks(xtics,['2^4','2^5','2^6','2^7'])
    plt.title('Norm')
    fig_norm.savefig(out_dir+'norm.pdf')
    
    

plot2d('./out/2D/FQM/')
plot2d('./out/2D/FQV/')
plot2d('./out/2D/Func2/')
plot2d('./out/2D/Func3/')

plot1d('./out/1D/FQM/',1)
plot1d('./out/1D/FQV/',1)
plot1d('./out/1D/Func2/',2)
plot1d('./out/1D/Func3/',3)
